import os

GROUPS = [
"01_Alkane",
"02_Alkene",
"03_Alkyne",
"04_Benzene",
"05_AlcoholEther",
"06_Amine",
"07_AldehydeKetone",
"08_Acids",
"09_Amides",
"10_AmineOxide",
"11_Ring3",
"12_Ring4",
"13_Ring5",
"14_Ring6",
"15_Chloro",
"16_Fluorin",
"17_Cynate",
"18_Gases",
"19_Nitrile",
"20_Nitro",
"21_Phasphazene",
"22_Silanes",
"23_SiliconCarbon",
"24_Siloxanes",
"25_Sulfur",
"26_Unhydrates",
"27_AromaticHalids",
"28_Ions",
#"B01_C",
#"B02_F",
#"B03_Cl",
#"B04_Br",
#"B05_O_Alcohols",
#"B06_N_Amine",
#"B07_N#C",
#"B08_N=C",
#"B09_N=N",
#"B10_NO2",
#"B11_Ethers",
#"B12_Amides",
#"B13_Amine_rings",
#"B14_Carbonyl",
#"B15_Peroxides",
#"B16_ORings",
#"B17_NOmixed",
#"B18_NOrings",
#"B19_rings",
#"B20_NX4",
#"B21_S",
#"B22_Smixed",
#"B23_SRings",
#"B24_Sulfonate",
#"B25_Si",
#"B26_SiO",
#"B27_Silicons",
#"B28_SiN",
#"B29_P",
#"B30_Se",
#"B31_Add",
#"BMB1",
#"BMB2",
#"BMB3",
#"BMB4",
#"BMB5",
#"BMB6",
#"BMB7",
#"BMB8",
#"BMB9",
]

DFF_ROOT = "E:\work\DFF 7"
DFI_NAME = "fit"
INIT_PPF = "initial.ppf"
INPUT_PPF = "input.ppf"
OUTPUT_PPF = "output.ppf"
DFFFIT_EXE = "dfffit"

class Parameter:
    def __init__(self, value = 0.0, fixed = False):
        self.value = value
        self.fixed = fixed

    def __str__(self):
        result = str(self.value)
        if (self.fixed):
            result += "*"
        return result

    def fix(self, fixed = True):
        self.fixed = fixed

def parseParameter(line):
    line = line.strip() # Remove white spaces
    fixed = False
    if line.endswith("*"):
        fixed = True
        line = line[:-1]
    return Parameter(float(line), fixed)

Parameter.parse = parseParameter

class ForceFieldItem:
    def __init__(self):
        self.term = "NONE"
        self.atoms = []
        self.pars = []
        self.tag = ""

    def __str__(self):
        result = self.term + ": " + ", ".join(self.atoms) + ": "
        for par in self.pars:
            result += str(par) + ", "
        result = result[:-2]
        if self.tag.strip() != "":
            result += ": " + self.tag
        return result

    def __eq__(self, item2):
        if self.term == item2.term:
            if len(self.atoms) == len(item2.atoms):
                for i in range(len(self.atoms)):
                    if self.atoms[i] != item2.atoms[i]:
                        return False
                return True
        return False

    def __ne__(self, item2):
        return not self.__eq__(item2)
                
    def fix(self, fixed = True):
        for i in self.pars:
            i.fix(fixed)

def parseForceFieldItem(line):
    line = line.strip()
    tokens = line.split(":", 3)
    result = ForceFieldItem()
    if len(tokens) > 0:
        result.term = tokens[0]
    if len(tokens) > 1:
        result.atoms = [i.strip() for i in tokens[1].split(",")]
    if len(tokens) > 2:
        result.pars = [parseParameter(i) for i in tokens[2].split(",")]
    if len(tokens) > 3:
        result.tag = tokens[3].strip()
    return result

ForceFieldItem.parse = parseForceFieldItem

class ForceField:
    def __init__(self):
        self.type = "NONE"
        self.data = []

    def fixAll(self):
        for i in self.data:
            i.fix(True)

    def write(self, file):
        file.write("#DFF:PPF\n#PROTOCOL = " + self.type+"\n")
        for item in self.data:
            file.write(str(item)+"\n")

    def index(self, item):
        for i in range(len(self.data)):
            if self.data[i] == item:
                return i
        return -1

    def __len__(self):
        return len(self.data)

def readForceField(file):
    result = ForceField()
    for line in file:
        line = line.strip()
        if line == "":
            break
        elif line.startswith("#DFF:PPF"):
            pass
        elif line.startswith("#PROTOCOL"):
            result.type = line.split('=')[1].strip()
        else:
            item = parseForceFieldItem(line)
            if item.term != "E0":
                result.data.append(item)
    return result
            
ForceField.read = readForceField

class EquivalenceTable:
    def __init__(self):
        self.data = { }

    def write(self, file):
        file.write("#AAT:\tNonbond\tBINC\tBond\tAngle/C\tAngle/S\tTorsion/C\tTorsion/S\tOOPA/C\tOOPA/S\n")
        for v in self.data.values():
            file.write(v[0] + ": " + "\t".join(v[1:]) + "\n")

    def fixAll(self):
        pass

    def __len__(self):
        return len(self.data)

def readEquivalenceTable(file):
    result = EquivalenceTable()
    for line in file:
        line = line.strip()
        if len(line) == 0:
            break
        elif line[0] != '#':
            tokens = line.split(": \t\n\r")
            tokens = [i.strip() for i in tokens]
            result.data[tokens[0]] = tokens
    return result
    
def fit(folder):
    print "####  Fitting", folder
    command = DFFFIT_EXE + ' "' + folder + '/' + DFI_NAME + '"'
    exe = os.popen(command, "r")
    for line in exe:
        print line.strip()
    exe.close()

def readLastOutputPPF(folder):
    ppfname = os.path.join(folder, OUTPUT_PPF)
    file = open(ppfname, "rt")
    ff = readForceField(file)
    file.close()
    return ff

def getEqtFileName(filename):
    index = filename.rfind(".")
    if index != -1:
        eqtname = filename[0:index]
    eqtname += ".eqt"
    return eqtname
    
def readLastOutputEQT(folder):
    eqtname = getEqtFileName(OUTPUT_PPF)
    eqtname = os.path.join(folder, eqtname)
    file = open(eqtname, "rt")
    eqt = readEquivalenceTable(file)
    file.close()
    return eqt

def readLastOutput(lastFolder):
    ff = readLastOutputPPF(lastFolder)
    eqt = readLastOutputEQT(lastFolder)
    return (ff, eqt)

def saveInput(folder, ff, eqt):
    ppfname = os.path.join(folder, INPUT_PPF)
    ppffile = open(ppfname, "wt")
    if ppffile:
        ff.write(ppffile)
        ppffile.close()

    eqtname = getEqtFileName(INPUT_PPF)
    eqtname = os.path.join(folder, eqtname)
    eqtfile = open(eqtname, "wt")
    if eqtfile:
        eqt.write(eqtfile)
        eqtfile.close()


def readInitPPF(folder):
    ppfname = os.path.join(folder, INIT_PPF)
    try:
        file = open(ppfname, "rt")
        ff = readForceField(file)
        file.close()
        return ff
    except IOError:
        return ForceField()
    
def readInitEQT(folder):
    eqtname = getEqtFileName(INIT_PPF)
    eqtname = os.path.join(folder, eqtname)
    print "$$$$$$$$", eqtname
    try:
        file = open(eqtname, "rt")
        print "$$$$$$$$", eqtname
        eqt = readEquivalenceTable(file)
        file.close()
        return eqt
    except IOError:
        return EquivalenceTable()

def readInitParameter(folder):
    ff = readInitPPF(folder)
    eqt = readInitEQT(folder)
    return (ff, eqt)

def mergeForceField(ff, ffI):
    if len(ffI) > 0:
        if ff.type != ffI.type:
            raise "Different Force Field Type"
        for i in ffI.data:
            if ff.index(i) != -1:
                raise "The item " + str(i) + " was fixed in input force field!"
            else:
                ff.data.append(i)
    return ff

def mergeEquivalenceTable(eqt, eqtI):
    for k, v in eqtI.data.iteritems():
        if k in eqt.data:
            raise "The EQT item " + k + " was fixed in input force field!"
        else:
            eqt.data[k] = v
    return eqt

def makeDFI(folder, template, fftype):
    global DFI_NAME, INIT_PPF, INPUT_PPF, OUTPUT_PPF
    DFI_NAME = "fit" + fftype
    INIT_PPF = "initial" + fftype + ".ppf"
    INPUT_PPF = "input" + fftype + ".ppf"
    OUTPUT_PPF = "output" + fftype + ".ppf"

    qmds = ""
    mols = ""

    for i in os.listdir("../Data/" + folder):
        if i.endswith(".qmd"):
            qmds += "  ../Data/" + folder + "/" + i + "\n"
            mols += "  ../Data/" + folder + "/" + i[:-4] + ".msd\n"
    qmds = qmds.strip()
    mols = mols.strip()

    template = template[:]
    template = template.replace("%DFF_ROOT%", DFF_ROOT)
    template = template.replace("%DATA_FILES%", qmds)
    template = template.replace("%MSD_FILES%", mols)
    template = template.replace("%FF_TYPE%", fftype)
    template = template.replace("%INPUT_FF%", folder + "/" + INPUT_PPF)
    template = template.replace("%OUTPUT_FF%", folder + "/" + OUTPUT_PPF)
    file = open(folder + "/" + DFI_NAME + ".dfi", "wt")
    file.write(template)
    file.close()

def autofit():
    fit(GROUPS[0])
    lastFolder = GROUPS[0]
    for folder in GROUPS[1:]:
        ff, eqt = readLastOutput(lastFolder)
        ff.fixAll()
        eqt.fixAll()

        ffI, eqtI = readInitParameter(folder)
        ff = mergeForceField(ff, ffI)
        eqt = mergeEquivalenceTable(eqt, eqtI)
        saveInput(folder, ff, eqt)
        fit(folder)
        lastFolder = folder

def autofitAll(type):
    file = open("fit.dfi", "rt")
    template = file.read()
    file.close()
    if not os.path.isdir("Fit" + type):
        os.mkdir("Fit" + type)
    os.chdir("Fit" + type)
    for i in GROUPS:
        if not os.path.isdir(i):
            os.mkdir(i)
        makeDFI(i, template, type)
    print DFI_NAME        
    autofit()
    os.chdir("..")


if __name__ == '__main__':
    autofitAll("AMBER")
    
